{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "This a notebook that inspects the results of a WarpX simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import statements\n",
    "import yt ; yt.funcs.mylog.setLevel(50)\n",
    "import numpy as np\n",
    "import scipy.constants as scc\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read data in the simulation frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Instruction\n",
    "\n",
    "Enter the path of the data you wish to visualize below. Then execute the cells one by one, by selecting them with your mouse and typing `Shift + Enter`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diag_name = 'diag' # E.g., diagnostics.diags_names = diag\n",
    "iteration = 0\n",
    "plotfile = './diags/{}{:05d}'.format(diag_name, iteration)\n",
    "field = 'Ex'\n",
    "species = 'electron'\n",
    "ds = yt.load( plotfile ) # Load the plotfile\n",
    "# ds.field_list # Print all available quantities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot data with yt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sl = yt.SlicePlot(ds, 2, field, aspect=.2) # Create a sliceplot object\n",
    "sl.annotate_particles(width=(10.e-6, 'm'), p_size=2, ptype=species, col='black')\n",
    "sl.annotate_grids() # Show grids\n",
    "sl.show() # Show the plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Store quantities in numpy arrays, and plot with matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get field quantities\n",
    "all_data_level_0 = ds.covering_grid(level=0,left_edge=ds.domain_left_edge, dims=ds.domain_dimensions)\n",
    "Bx = all_data_level_0['boxlib', field].v.squeeze()\n",
    "Dx = ds.domain_width/ds.domain_dimensions\n",
    "extent = [ds.domain_left_edge[ds.dimensionality-1], ds.domain_right_edge[ds.dimensionality-1],\n",
    "          ds.domain_left_edge[0], ds.domain_right_edge[0] ]\n",
    "\n",
    "# Get particle quantities\n",
    "ad = ds.all_data()\n",
    "x = ad[species, 'particle_position_x'].v\n",
    "z = ad[species, 'particle_position_z'].v\n",
    "\n",
    "# Plot image\n",
    "plt.figure()\n",
    "plt.imshow(Bx[:,Bx.shape[1]//2,:], extent=extent, aspect='auto')\n",
    "plt.scatter(z,x,s=.1,c='k')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read data back-transformed to the lab frame when the simulation runs in the boosted frame (example: 2D run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read_raw_data.py is located in warpx/Tools.\n",
    "import os, glob\n",
    "import read_raw_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration = 1\n",
    "\n",
    "snapshot = './lab_frame_data/' + 'snapshot' + str(iteration).zfill(5)\n",
    "header   = './lab_frame_data/Header'\n",
    "allrd, info = read_raw_data.read_lab_snapshot(snapshot, header) # Read field data\n",
    "F = allrd[field]\n",
    "print( \"Available info: \", list(info.keys()) )\n",
    "print(\"Available fields: \", info['field_names'])\n",
    "nx = info['nx']\n",
    "nz = info['nz']\n",
    "x = info['x']\n",
    "z = info['z']\n",
    "xbo  = read_raw_data.get_particle_field(snapshot, species, 'x') # Read particle data\n",
    "ybo  = read_raw_data.get_particle_field(snapshot, species, 'y')\n",
    "zbo  = read_raw_data.get_particle_field(snapshot, species, 'z')\n",
    "uzbo = read_raw_data.get_particle_field(snapshot, species, 'uz')\n",
    "\n",
    "plt.figure(figsize=(6, 3))\n",
    "extent = np.array([info['zmin'], info['zmax'], info['xmin'], info['xmax']])\n",
    "plt.imshow(F[:,F.shape[1]//2,:], aspect='auto', extent=extent, cmap='seismic')\n",
    "plt.colorbar()\n",
    "plt.plot(zbo, xbo, 'g.', markersize=1.)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read back-transformed data with hdf5 format (example: 3D run)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import matplotlib.pyplot as plt\n",
    "f = h5py.File('HDF5_lab_frame_data/snapshot00003', 'r')\n",
    "print( list(f.keys()) )\n",
    "# plt.figure()\n",
    "plt.imshow(f['Ey'][:,,:])"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  },
  "widgets": {
   "state": {
    "11d243e9f5074fe1b115949d174d59de": {
     "views": [
      {
       "cell_index": 6
      }
     ]
    }
   },
   "version": "1.2.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
